import os
import base64
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from OpenSSL import crypto

def print_header(title):
    """美观的标题打印"""
    print(f"\n{'='*40}")
    print(f"=== {title}")
    print(f"{'='*40}\n")

def print_hex(data, prefix="", length=16):
    """以十六进制格式打印二进制数据"""
    hex_str = data.hex()[:length*2]
    truncated = "(truncated)" if len(data) > length else ""
    print(f"{prefix}{hex_str} {truncated}[Total {len(data)} bytes]")

def save_encrypted_key(encrypted_data, filename="encrypted_session_key.bin"):
    """保存加密后的会话密钥到文件"""
    try:
        with open(filename, "wb") as f:
            f.write(encrypted_data)
        print(f"\n[操作] 加密会话密钥已保存到: {os.path.abspath(filename)}")
        print_hex(encrypted_data, prefix="  加密密钥内容: ", length=12)
        print(f"  Base64编码预览: {base64.b64encode(encrypted_data[:16]).decode()}...")
    except Exception as e:
        print(f"[错误] 保存失败: {str(e)}")
        exit(1)

# ----------------- 主流程 -----------------
if __name__ == "__main__":
    print_header("0. 初始化物理机")
    
    # 加载虚拟机证书
    print_header("1. 加载虚拟机证书")
    try:
        with open("vm.crt", "rb") as f:
            vm_cert = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
        vm_pubkey = vm_cert.get_pubkey().to_cryptography_key()
        
        # 验证公钥类型
        print("[证书信息验证]")
        print(f"  证书主题: {vm_cert.get_subject().CN}")
        print(f"  公钥类型: {type(vm_pubkey).__name__}")
        print(f"  密钥长度: {vm_pubkey.key_size} bits")
        
    except FileNotFoundError:
        print("[错误] 找不到 vm.crt，请先运行证书生成脚本！")
        exit(1)
    except Exception as e:
        print(f"[错误] 证书加载失败: {str(e)}")
        exit(1)

    # 生成会话密钥
    print_header("2. 生成会话密钥")
    session_key = os.urandom(32)  # AES-256
    print_hex(session_key, prefix="  原始会话密钥: ", length=8)
    print(f"  Base64完整编码: {base64.b64encode(session_key).decode()}")

    # RSA加密会话密钥
    print_header("3. 加密会话密钥")
    try:
        encrypted_key = vm_pubkey.encrypt(
            session_key,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
    except Exception as e:
        print(f"[加密失败] 错误类型: {type(e).__name__}, 详情: {str(e)}")
        exit(1)

    # 验证加密结果
    print_header("4. 验证加密结果")
    expected_length = vm_pubkey.key_size // 8  # 2048位密钥对应256字节
    print(f"  理论密文长度: {expected_length} bytes")
    print(f"  实际密文长度: {len(encrypted_key)} bytes")
    
    if len(encrypted_key) != expected_length:
        print("[致命错误] 加密后的会话密钥长度不符合密钥长度！")
        exit(1)
    else:
        print("[验证通过] 加密密钥长度符合要求")

    # 保存加密后的密钥
    print_header("5. 保存加密密钥")
    save_encrypted_key(encrypted_key)

    print("\n" + "="*40)
    print("=== 物理机端操作完成 ===")
    print("="*40)